import base64
import sys
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Random import get_random_bytes


def base64Encode(binData):
    bs64 = base64.b64encode(binData).decode("utf-8")
    return bs64.replace("/", "_")


def base64Decode(bs64):
    try:
        bs64 = bs64.replace("_", "/")
        return base64.b64decode(bs64.encode("utf-8"))
    except Exception as e:
        print(f"解密出错：{str(e)}")
        sys.exit(1)


def aesEncrypt(text, pwd):
    salt = get_random_bytes(16)
    key = PBKDF2(
        password=pwd,
        salt=salt,
        dkLen=32,
        count=100000
    )
    cipher = AES.new(key, AES.MODE_GCM)
    textBin = text.encode("utf-8")
    encryptedBin, tag = cipher.encrypt_and_digest(textBin)
    return salt + cipher.nonce + tag + encryptedBin


def aesDecrypt(encryptedBin, pwd):
    try:
        salt = encryptedBin[:16]
        nonce = encryptedBin[16:32]
        tag = encryptedBin[32:48]
        encryptedData = encryptedBin[48:]

        key = PBKDF2(
            password=pwd,
            salt=salt,
            dkLen=32,
            count=100000
        )
        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
        decryptedBin = cipher.decrypt_and_verify(encryptedData, tag)
        return decryptedBin.decode("utf-8")
    except Exception as e:
        print(f"解密出错：{str(e)}")
        sys.exit(1)
